// EFI_FWInfo.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include <stdint.h>
#include <windows.h>

const char* EFI_GLOBAL_VARIABLE = "{8BE4DF61-93CA-11D2-AA0D-00E098032B8C}";
//const char* EFI_GLOBAL_VARIABLE = "{00000000-0000-0000-0000-000000000000}";

void PrintError(DWORD errorCode) {
    LPVOID lpMsgBuf;

    FormatMessage(
        FORMAT_MESSAGE_ALLOCATE_BUFFER |
        FORMAT_MESSAGE_FROM_SYSTEM |
        FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL,
        errorCode,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPTSTR)&lpMsgBuf,
        0,
        NULL);

    printf("Error code was: %ld\n", errorCode);
    MessageBox(0,(LPCWSTR)lpMsgBuf, L"Error",48);
    LocalFree(lpMsgBuf);
}

int read_efi_variable(const char* name, uint16_t** data) {
    // Windows does not allocate memory for it so we have to do it
    uint16_t* res = (uint16_t *)malloc(sizeof(uint16_t)* 1024);
    int data_size = GetFirmwareEnvironmentVariableA(name, EFI_GLOBAL_VARIABLE, res, sizeof(uint16_t)* 1024);
    if (data_size) {
        *data = res;
        return data_size / 2;
    }
    else {
        PrintError(GetLastError());
        exit(GetLastError());
    }
}

int set_efi_variable(const char* name, uint16_t num) {
    int res = SetFirmwareEnvironmentVariableA(name, EFI_GLOBAL_VARIABLE, (PVOID)&num, sizeof(uint16_t));
    if (res) {
        return res;
    }
    else {
        PrintError(GetLastError());
        exit(GetLastError());
    }
}

int set_boot_next(uint16_t num) {
    return set_efi_variable("BootNext", num);
}

int boot_entry_list(uint16_t** list) {
    uint16_t *data = NULL;
    int length = read_efi_variable("BootOrder", &data);
    *list = data;
    return length;
}

char* boot_entry_name(uint16_t id) {
    char buffer[9];
    uint16_t *data = NULL;
    int length, i;
    char* name = (char *)malloc(64 * sizeof(char));
    sprintf_s(buffer, 9, "Boot%04x", id);
    length = read_efi_variable(buffer, &data);
    for (i = 3; i < length; i++) {
        if (data[i] == 0) break;
        name[i - 3] = data[i];
    }
    name[i - 3] = 0;
    return name;
}

void ObtainPrivileges(LPCTSTR privilege) {
    HANDLE hToken;
    TOKEN_PRIVILEGES tkp;
    BOOL res;
    DWORD error;
    // Obtain required privileges
    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
        printf("OpenProcessToken failed!\n");
        PrintError(GetLastError());
        exit(GetLastError());
    }

    res = LookupPrivilegeValue(NULL, privilege, &tkp.Privileges[0].Luid);
    if (!res) {
        printf("LookupPrivilegeValue failed!\n");
        PrintError(GetLastError());
        exit(GetLastError());
    }
    tkp.PrivilegeCount = 1;
    tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
    AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, (PTOKEN_PRIVILEGES)NULL, 0);

    error = GetLastError();
    if (error != ERROR_SUCCESS) {
        printf("AdjustTokenPrivileges failed\n");
        PrintError(error);
        exit(error);
    }

}

void Usage(void)
{
	fprintf(stderr, "NextBoot for Windows with EFI BIOS    Rev:1.0.0");
	fprintf(stderr, "\nnextBoot.exe -n <ID> | -d");
	fprintf(stderr, "\nOption");
	fprintf(stderr, "\n\tn	next boot option");
	fprintf(stderr, "\n\tID	next boot device");
	fprintf(stderr, "\n\td	dump boot devices");
	fprintf(stderr, "\n\nReturn code");
	fprintf(stderr, "\n\t0	Pass");
	fprintf(stderr, "\n\t1	Fail");
	exit(1);
}

int main(int argc, char* argv[])
{
	uint16_t *data;
    int l, i, c;
    data = NULL;

	ObtainPrivileges(SE_SYSTEM_ENVIRONMENT_NAME);

	if (argc == 1)
		Usage();

	if (argc == 2 && strcmp(argv[1],"/?") == 0 ){
		Usage();
	}

	// Get boot data...
	l = boot_entry_list(&data);
	if (l < 1){
		printf("Error: No any Boot device found or NOT EFI BIOS\n");
		exit(1);
	}

	if (argc == 2 && strcmp(argv[1], "-d") == 0){
		l = boot_entry_list(&data);
		for (i = 0; i < l; i++) {
			printf("%d: %s\n", i, boot_entry_name(data[i]));
		}
	}else if (argc ==3 && strcmp(argv[1], "-n") ==0){
		//printf("next boot choice: ");
		if(argv[2][0] < '0' ||  argv[2][0] > '9'){
			printf("Error: Option ID is wrong(0~9 is accepted)\n");
			exit(1);
		}

		c=atoi(argv[2]);
		//scanf("%d", &c);
		if (c < 0 || c >= l) {
			printf("Error: Invalid boot ID(No this device)\n");
			exit(1);
		}
		set_boot_next(data[c]);
		printf("Set next boot target to: %s\n", boot_entry_name(data[c]));
	}else{
		Usage();
	}

	return 0;
}

